/*
多线程单卡测试场景，3个线程使用最后1张卡
*/
#include <cassert>
#include <cstdlib>
#include <hip/hip_runtime.h>
#include <iostream>
#include <sys/wait.h>
#include <thread>
#include <unistd.h>
#include <vector>

#define HIPCHECK(call)                                                            \
  do {                                                                         \
    hipError_t err = call;                                                     \
    if (err != hipSuccess) {                                                   \
      std::cerr << "HIP Error [" << __FILE__ << ":" << __LINE__               \
                << "] Code=" << err << "(" << hipGetErrorString(err)           \
                << ") PID=" << getpid()                                        \
                << " TID=" << std::this_thread::get_id() << std::endl;         \
      std::abort();                                                            \
    }                                                                          \
  } while (0)

__global__ void computeKernel(float *a, float *b, float *c, int n) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < n) {
    c[idx] = sinf(a[idx]) * 2.0f + cosf(b[idx]) * 0.5f;
  }
}

void verifyResults(const float *result, const float *expected, int size) {
  const float epsilon = 1e-4f;
  for (int i = 0; i < size; ++i) {
    if (fabs(result[i] - expected[i]) > epsilon) {
      std::cerr << "Result check failure" << std::endl;
      // std::abort();
    }
  }
}

void gpuWorker(int deviceId, int dataSize) {
  // 设置设备上下文
  HIPCHECK(hipSetDevice(deviceId));

  // 验证设备设置
  int currentDevice;
  HIPCHECK(hipGetDevice(&currentDevice));
  assert(currentDevice == deviceId);

  const size_t bytes = dataSize * sizeof(float);

  // 使用分页锁定内存
  float *hA, *hB, *hC, *hVerify;
  HIPCHECK(hipHostMalloc(&hA, bytes));
  HIPCHECK(hipHostMalloc(&hB, bytes));
  HIPCHECK(hipHostMalloc(&hC, bytes));
  HIPCHECK(hipHostMalloc(&hVerify, bytes));

  // 初始化数据
  for (int i = 0; i < dataSize; ++i) {
    hA[i] = static_cast<float>(i) * 0.1f;
    hB[i] = static_cast<float>(dataSize - i) * 0.1f;
    hVerify[i] = sinf(hA[i]) * 2.0f + cosf(hB[i]) * 0.5f;
  }

  // 设备内存分配
  float *dA, *dB, *dC;
  HIPCHECK(hipMalloc(&dA, bytes));
  HIPCHECK(hipMalloc(&dB, bytes));
  HIPCHECK(hipMalloc(&dC, bytes));

  // 创建流
  hipStream_t stream;
  HIPCHECK(hipStreamCreate(&stream));

  // 异步数据传输
  HIPCHECK(hipMemcpyAsync(dA, hA, bytes, hipMemcpyHostToDevice, stream));
  HIPCHECK(hipMemcpyAsync(dB, hB, bytes, hipMemcpyHostToDevice, stream));

  // 执行核函数
  const int blockSize = 256;
  const int gridSize = (dataSize + blockSize - 1) / blockSize;
  computeKernel<<<gridSize, blockSize, 0, stream>>>(dA, dB, dC, dataSize);

  // 错误检查
  HIPCHECK(hipGetLastError());
  HIPCHECK(hipStreamSynchronize(stream));

  // 回传结果
  HIPCHECK(hipMemcpyAsync(hC, dC, bytes, hipMemcpyDeviceToHost, stream));
  HIPCHECK(hipStreamSynchronize(stream));

  // 验证结果
  verifyResults(hC, hVerify, dataSize);

  // 资源清理
  HIPCHECK(hipFree(dA));
  HIPCHECK(hipFree(dB));
  HIPCHECK(hipFree(dC));
  HIPCHECK(hipHostFree(hA));
  HIPCHECK(hipHostFree(hB));
  HIPCHECK(hipHostFree(hC));
  HIPCHECK(hipHostFree(hVerify));
  HIPCHECK(hipStreamDestroy(stream));

  std::cout << "[PID:" << getpid() << "][TID:" << std::this_thread::get_id()
            << "] GPU " << deviceId << " validation passed\n";
}

int main() {
  int deviceCount = 0;
  HIPCHECK(hipGetDeviceCount(&deviceCount));
  std::cout << "check device num " << deviceCount << std::endl;

  if (deviceCount == 0) {
    std::cerr << "No DCU devices found." << std::endl;
    return 1;
  }

  const int dataSize = 1 << 20; // 1M elements
  std::vector<std::thread> threads;

//  for (int i = 0; i < deviceCount; ++i) {
//    threads.emplace_back(gpuWorker, i, dataSize);
//  }
  int last_device_id = deviceCount - 1;
  for (int i = 0; i < 3; ++i) {  // 设置3个线程使用最后一张卡
    threads.emplace_back(gpuWorker, last_device_id, dataSize);
  }

  // 等待所有线程结束
  for (auto &thread : threads) {
    thread.join();
  }

  std::cout << "All processes finished." << std::endl;
  return 0;
}
